import sys
import os
import time
import tensorflow as tf
import numpy as np
from modelscope.models.audio.tts import SambertHifigan
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

#model_dir = os.path.abspath("./pretrain_yangmi")
model_dir = os.path.abspath(sys.argv[1])

custom_infer_abs = {
    'voice_name':
    'F7',
    'am_ckpt':
    os.path.join(model_dir, 'tmp_am', 'ckpt'),
    'am_config':
    os.path.join(model_dir, 'tmp_am', 'config.yaml'),
    'voc_ckpt':
    os.path.join(model_dir, 'orig_model', 'basemodel_16k', 'hifigan', 'ckpt'),
    'voc_config':
    os.path.join(model_dir, 'orig_model', 'basemodel_16k', 'hifigan',
             'config.yaml'),
    'audio_config':
    os.path.join(model_dir, 'data', 'audio_config.yaml'),
    'se_file':
    os.path.join(model_dir, 'data', 'se', 'se.npy')
}
kwargs = {'custom_ckpt': custom_infer_abs}

model_id = SambertHifigan(os.path.join(model_dir, "orig_model"), **kwargs)

inference = pipeline(task=Tasks.text_to_speech, model=model_id, device='cpu')
srctxt = sys.argv[2]
dstpath = sys.argv[3]
with open(srctxt, 'r') as f:
   alltxt = f.readlines() 
for txtline in alltxt:
    txtline=txtline.strip().split()
    start = time.time()
    output = inference(input=txtline[1])
    print('spend time {}'.format(time.time()-start))
    wav = output[OutputKeys.OUTPUT_WAV]
    with open(os.path.join(dstpath,txtline[0]+".wav"), 'wb') as f:
        f.write(wav)

print('am para')
am_model_para = model_id.voices['F7'].am.state_dict()
for k,v in am_model_para.items():
     print("{:20s} {}".format(k, v.shape))

print('voc para')
voc_model_para = model_id.voices['F7'].voc_model.state_dict()
for k,v in voc_model_para.items():
     print("{:20s} {}".format(k, v.shape))

def binarize(mel, threshold=0.6):
    mask = mel[:,:,-1] < threshold
    vflag = tf.where(mask,tf.zeros([1]),tf.ones([1]))
    res_mel = tf.concat([mel[:,:,:-1],tf.expand_dims(vflag,axis=-1)],axis=-1)
    return res_mel

def voc_infer(post_mel):

    class ResidualBlock(tf.keras.layers.Layer):
        def __init__(self, channels,kernel_size, dilation=[1,3,5,7],**kwargs):
            super(ResidualBlock, self).__init__(**kwargs)

            self.convs1 = []
            for li in range(4):
                cv1 = tf.keras.Sequential()
                cv1.add(tf.keras.layers.LeakyReLU(alpha=0.1))
                cv1.add(tf.keras.layers.Conv1D(channels,kernel_size,strides=1,padding='same',dilation_rate=dilation[li]))
                self.convs1.append(cv1)

            self.convs2 = []
            for li in range(4):
                cv2 = tf.keras.Sequential()
                cv2.add(tf.keras.layers.LeakyReLU(alpha=0.1))
                cv2.add(tf.keras.layers.Conv1D(channels,kernel_size,strides=1,padding='same',dilation_rate=1))
                self.convs2.append(cv2)

        def call(self,x):
            for i in range(4):
                xt = self.convs1[i](x)
                xt = self.convs2[i](xt)
                x = xt + x
            return x

    class SourceModule(tf.keras.layers.Layer):
        def __init__(self,nb_harmonics, upsample_ratio, sampling_rate, alpha=0.1, sigma=0.003,**kwargs):
            super(SourceModule, self).__init__(**kwargs)
            self.nb_harmonics = nb_harmonics
            self.upsample_ratio = upsample_ratio
            self.sampling_rate = sampling_rate
            self.alpha = alpha
            self.sigma = sigma

            self.ffn = tf.keras.layers.Conv1D(1,1,padding='same',activation=tf.keras.activations.tanh)
            self.ups = tf.keras.layers.UpSampling1D(size=upsample_ratio)

        def call(self,pitch, uv):
            pitch_samples = self.ups(pitch)
            uv_samples = self.ups(uv)
            F_mat_a = []
            for i in range(self.nb_harmonics+1):
                F_mat_a.append(pitch_samples * (i + 1) / self.sampling_rate)
            F_mat = tf.concat(F_mat_a,axis=-1)
            theta_mat = 2 * np.pi * (tf.math.cumsum(F_mat, axis=1) % 1)
            u_dist = tf.random.uniform(shape=(1,1,self.nb_harmonics+1),minval=-np.pi, maxval=np.pi)
            phase_vec = tf.concat([tf.zeros([1,1,1]),u_dist[:,:,1:]],axis=-1)

            noise = tf.random.normal(shape=(1,tf.shape(pitch_samples)[1],self.nb_harmonics+1),mean=0.0,stddev=self.sigma)
            e_voice = self.alpha * tf.math.sin(theta_mat + phase_vec) + noise
            e_unvoice = self.alpha / 3 / self.sigma * noise
            e = e_voice * uv_samples + e_unvoice * (1 - uv_samples)
            return self.ffn(e)
    
    class UpsamModule(tf.keras.layers.Layer):
        def __init__(self,upsample_ratio,**kwargs):
            super(UpsamModule, self).__init__(**kwargs)
            self.upsample_ratio = upsample_ratio

        def call(self,x):
            return tf.repeat(x,repeats=[self.upsample_ratio],axis=1)

    class Generator(tf.keras.layers.Layer):
        def __init__(self, channels=256,kernel_size=7, upsample_kernal_sizes=[20, 11, 4, 4],upsample_scales=[10, 5, 2, 2],resblock_kernel_sizes=[3, 7, 11],**kwargs):
            super(Generator, self).__init__(**kwargs)

            self.num_upsamples = len(upsample_kernal_sizes)
            self.num_kernels = len(resblock_kernel_sizes)
            self.transpose_upsamples = []
            for li in range(self.num_upsamples):
                trans_upsamples = tf.keras.Sequential()
                trans_upsamples.add(tf.keras.layers.LeakyReLU(alpha=0.1))
                trans_upsamples.add(tf.keras.layers.Conv1DTranspose(channels//(2**(li+1)),upsample_kernal_sizes[li],upsample_scales[li],padding='same'))
                self.transpose_upsamples.append(trans_upsamples)


            self.repeat_upsamples = []
            for li in range(self.num_upsamples):
                rep_upsamples = tf.keras.Sequential()
                #rep_upsamples.add(tf.keras.layers.UpSampling1D(size=upsample_scales[li]))
                rep_upsamples.add(UpsamModule(upsample_scales[li]))
                rep_upsamples.add(tf.keras.layers.LeakyReLU(alpha=0.1))
                rep_upsamples.add(tf.keras.layers.Conv1D(channels//(2**(li+1)),kernel_size,strides=1,padding='same'))
                self.repeat_upsamples.append(rep_upsamples)

            self.conv_blocks = []
            for li in range(self.num_upsamples):
                for lj in range(self.num_kernels):
                    self.conv_blocks.append(ResidualBlock(channels//(2**(li+1)),resblock_kernel_sizes[lj]))

            self.conv_pre = tf.keras.layers.Conv1D(channels,kernel_size,padding='same')
            self.conv_post = tf.keras.layers.Conv1D(1,kernel_size,padding='same')
            self.source_module = SourceModule(7,np.cumprod(upsample_scales)[-1],16000)

            downsample_rates = [1] + upsample_scales[::-1][:-1]
            downsample_cum_rates = np.cumprod(downsample_rates)[::-1].tolist()
            self.source_downs = []
            for li in range(3):
                self.source_downs.append(tf.keras.layers.Conv1D(channels//(2**(li+1)),kernel_size=2*downsample_cum_rates[li],strides=downsample_cum_rates[li],padding='same'))
            self.source_downs.append(tf.keras.layers.Conv1D(channels//(2**(3+1)),1,1,padding='same'))
            self.leac = tf.keras.layers.LeakyReLU(alpha=0.01)

        def call(self, x):
            x = binarize(x)
            mel = x[:, :, :-2]
            pitch = x[:, :, -2:-1]
            uv = x[:, :, -1:]
            excitation = self.source_module(pitch, uv)

            x = self.conv_pre(mel)

            for i in range(self.num_upsamples):
                x = tf.math.sin(x) + x
                rep = self.repeat_upsamples[i](x)
                up = self.transpose_upsamples[i](x)
                e = self.source_downs[i](excitation)
                x = rep + e + up[:, :tf.shape(rep)[1], :]

                xs = tf.zeros([1])
                for j in range(self.num_kernels):
                    xs += self.conv_blocks[i * self.num_kernels + j](x)
                x = xs / self.num_kernels
            x = self.leac(x)
            x = self.conv_post(x)
            x = tf.keras.activations.tanh(x)
            return x
        
    mel_input = tf.keras.Input(shape=(None,82),batch_size=1,name='mel_in')
    pcm_out = Generator(name='hifigan')(mel_input)
    hifi_model = tf.keras.Model((mel_input),(pcm_out))
    hifi_model.summary()
    hifi_model_layer = hifi_model.get_layer(name='hifigan')
    #weights = hifi_model.weights
    #for weight in weights:
    #   print('weight {}'.format(weight.name))
    #   print('weight {}'.format(weight.shape))
    #   #print('weight {}'.format(weight))
    hifi_para = []
    for i in range(4):
        trans_conv_w = np.transpose(voc_model_para["transpose_upsamples.{}.1.deconv.weight".format(i)].cpu().numpy(),(2,1,0))
        trans_conv_b = voc_model_para["transpose_upsamples.{}.1.deconv.bias".format(i)].cpu().numpy()
        hifi_para.extend([trans_conv_w,trans_conv_b])
    for i in range(4):
        rep_conv_w = np.transpose(voc_model_para["repeat_upsamples.{}.2.conv1d.weight".format(i)].cpu().numpy(),(2,1,0))
        rep_conv_b = voc_model_para["repeat_upsamples.{}.2.conv1d.bias".format(i)].cpu().numpy()
        hifi_para.extend([rep_conv_w,rep_conv_b])
    for i in range(12):
        for j in range(4):
            conv1_conv_w = np.transpose(voc_model_para["conv_blocks.{}.convs1.{}.conv1d.weight".format(i,j)].cpu().numpy(),(2,1,0))
            conv1_conv_b = voc_model_para["conv_blocks.{}.convs1.{}.conv1d.bias".format(i,j)].cpu().numpy()
            hifi_para.extend([conv1_conv_w,conv1_conv_b])
        for j in range(4):
            conv2_conv_w = np.transpose(voc_model_para["conv_blocks.{}.convs2.{}.conv1d.weight".format(i,j)].cpu().numpy(),(2,1,0))
            conv2_conv_b = voc_model_para["conv_blocks.{}.convs2.{}.conv1d.bias".format(i,j)].cpu().numpy()
            hifi_para.extend([conv2_conv_w,conv2_conv_b])

    pre_conv_w = np.transpose(voc_model_para["conv_pre.conv1d.weight"].cpu().numpy(),(2,1,0))
    pre_conv_b = voc_model_para["conv_pre.conv1d.bias"].cpu().numpy()
    post_conv_w = np.transpose(voc_model_para["conv_post.conv1d.weight"].cpu().numpy(),(2,1,0))
    post_conv_b = voc_model_para["conv_post.conv1d.bias"].cpu().numpy()
    source_conv_w = np.transpose(voc_model_para["source_module.ffn.0.weight"].cpu().numpy(),(2,1,0))
    source_conv_b = voc_model_para["source_module.ffn.0.bias"].cpu().numpy()
    hifi_para.extend([pre_conv_w,pre_conv_b,post_conv_w,post_conv_b,source_conv_w,source_conv_b])
    for i in range(4):
        sdown_conv_w = np.transpose(voc_model_para["source_downs.{}.conv1d.weight".format(i)].cpu().numpy(),(2,1,0))
        sdown_conv_b = voc_model_para["source_downs.{}.conv1d.bias".format(i)].cpu().numpy()
        hifi_para.extend([sdown_conv_w,sdown_conv_b])
    hifi_model_layer.set_weights(hifi_para)
    hifi_model.save('tfmodel/hifigan')

    #hifi_in = tf.expand_dims(tf.convert_to_tensor(post_mel.cpu().numpy()),axis=0)
    hifi_in = tf.expand_dims(tf.convert_to_tensor(post_mel),axis=0)
    hifi_out = hifi_model(hifi_in)
    pcm = (tf.reshape(hifi_out,(-1))*32768.0).numpy().astype(dtype=np.int16)
    import wave
    with wave.open("tfdebug.wav", 'wb') as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(16000)
        wf.writeframes(pcm.data)



#voc_infer(model_id.voices['F7'].tmp_vocinput)
